package com.lvyh.lightframe.log.thread;

import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;

/**
 * MDC transfer of asynchronous subtasks.
 * Customize thread pool to pass MDC of parent thread to child thread.
 *
 * @author lvyh 2021/03/20.
 */
public class MDCThreadPoolTaskExecutor extends ThreadPoolTaskExecutor {

    @Override
    public void execute(Runnable runnable) {
        Map<String, String> context = MDC.getCopyOfContextMap();//MDC content of parent thread.
        super.execute(() -> run(runnable, context));
    }

    /**
     * The delegate execution method of sub thread.
     */
    private void run(Runnable runnable, Map<String, String> context) {
        // Check whether it is null. If it is not checked, when the content is null (the parent thread has no MDC), a null pointer exception will be reported.
        if (null != context) {
            // Pass the MDC content of the parent thread to the child thread.
            MDC.setContextMap(context);
        }

        try {
            runnable.run();
        } finally {
            MDC.clear();
        }
    }
}
